# coding:utf-8
import operator
import matplotlib.pyplot as plt
import pandas as pd
import random
import numpy as np
from math import log


def calc_shannon_ent(data_set):
    num_entries = len(data_set)      # 传进来的数据的长度
    label_counts = {}                # 创建空字典，key为数据，value为数据出现的次数的字典
    for feat_vec in data_set:       # 循环得出数据
        current_label = feat_vec[-1]    # 数据
        if current_label not in label_counts.keys():  # 如果不在label_counts中，就添加
            label_counts[current_label] = 0
        label_counts[current_label] += 1            # 在的话，value+1
    shannon_ent = 0.0   #香农熵
    for key in label_counts:
        prob = float(label_counts[key])/num_entries     # 概率
        shannon_ent -= prob * log(prob, 2)    # 香农熵公式   H = -p(Xi) * log(prob, 2)   负的期望值乘期望值的对数
    return shannon_ent                        # 返回香农熵


def split_data_set(data_set, axis, value):
    ret_data_set = []                          # 由于python不考虑内存分配，对列表直接操作影响其生命周期，所以新建一个空列表
    for feat_vec in data_set:
        if feat_vec[axis] == value:
            reduce_feat_vec = feat_vec[:axis]
            reduce_feat_vec.extend(feat_vec[axis+1:])    # extend直接合并成一个列表，没有字list
            ret_data_set.append(reduce_feat_vec)           # 保持原有的list进行合并
    return ret_data_set                                   # 返回的是传入的axis和value相等的数组的[axis+1:]的值


def choose_best_feature2split(data_set):
    num_features = len(data_set[0]) - 1              # 每个数组内的属特征值的个数-1
    base_entropy = calc_shannon_ent(data_set)        # 基础的香农熵
    best_info_gain = 0.0                             # 要选择最高的信息增益
    best_feature = -1                                # 最好的特征的索引，-1表示不存在，还没开始选择
    for i in range(num_features):
        feat_list = [example[i] for example in data_set]        # 每个数组内的第i个值
        unique_vals = set(feat_list)                               # 去重，缺点连续值时，数据量太大
        new_entropy = 0.0                                                 # 新的香农熵
        for value in unique_vals:
            sub_data_set = split_data_set(data_set, i, value)
            prob = len(sub_data_set)/float(len(data_set))
            new_entropy += prob * calc_shannon_ent(sub_data_set)
        info_gain = (base_entropy - new_entropy)
        if info_gain > best_info_gain:
            best_info_gain = info_gain
            best_feature = i
    return best_feature


# 只能传一列label，也就是m*1矩阵
def majority_count(class_list):
    class_count = {}
    for vote in class_list:
        if vote not in class_count.keys():
            class_count[vote] = 0
        class_count[vote] = 1
    sorted_class_count = sorted(class_count.items(), key=operator.itemgetter(1), reverse=True)
    return sorted_class_count[0][0]      # 降序排列label出现的次数


def classify(input_tree, feat_labels, test_vec):
    first_str = first_str = list(input_tree.keys())[0]
    second_dict = input_tree[first_str]
    feat_index = feat_labels.index(first_str)
    for key in second_dict.keys():
        if test_vec[feat_index] == key:
            if type(second_dict[key]).__name__ == 'dict':
                class_label = classify(second_dict[key], feat_labels, test_vec)
            else:
                class_label = second_dict[key]
    return class_label


def store_tree(input_tree, file_name):
    import pickle
    fw = open(file_name, 'w')
    pickle.dump(input_tree, fw)
    fw.close()


def grab_tree(file_name):
    import pickle
    fr = open(file_name)
    return pickle.load(fr)


def create_tree(data_set, labels):
    class_list = [example[-1] for example in data_set]          # 取每行的最后一个元素
    # 出现class_list[0]的次数是否等于class_list的长度，
    if class_list.count(class_list[0]) == len(class_list):
        return class_list[0]
    if len(data_set[0]) == 1:
        return majority_count(class_list)
    best_feat = choose_best_feature2split(data_set)
    best_feat_label = labels[best_feat]
    my_tree = {best_feat_label: {}}
    del(labels[best_feat])
    feat_values = [example[best_feat] for example in data_set]
    unique_values = set(feat_values)
    for value in unique_values:
        sub_labels = labels[:]
        my_tree[best_feat_label][value] = create_tree(split_data_set(data_set, best_feat,value), sub_labels)
    return my_tree


my_data = [[1, 1, 'yes'], [1, 1, 'yes'], [1, 0, 'no'], [0, 0, 'no'], [0, 0, 'no']]
my_labels = ['no surfacing', 'flippers']
# print(calc_shannon_ent(my_data))
print(create_tree(my_data, my_labels))















